import pickle
import numpy as np
from configuration import DataConfig, TrainingConfig


class Data:
    def __init__(self, ):
        self.config = DataConfig()
        self.training_config = TrainingConfig()

        self.data_preprocess()

    def data_preprocess(self):
        """处理数据：提取图片名称、对应caption；建立字典；建立char_to_id、id_to_char；构建caption的vector"""
        # 从token file中读取image_names 和captions并加入开头结尾，存入list
        self.images_names = []
        self.captions = []
        token_lines = open(self.config.token_file, 'r').read().strip().split('\n')
        for i, row in enumerate(token_lines):
            row = row.split('\t')
            row[0] = row[0][:-2]
            self.images_names.append(row[0])
            self.captions.append('<start> ' + row[1] + ' <end>')

        # 建立字典，all_words:
        all_words = []
        for caption in self.captions:
            all_words.extend(caption.split())
        to_lower = lambda x: x.lower()
        all_words = list(set(map(to_lower, all_words)))
        all_words.append('<unk>')
        pickle.dump(all_words, open(self.config.vocab_file, 'wb'))

        # char to id &&& id to char
        char_to_id = {ch: i for i, ch in enumerate(all_words)}
        id_to_char = {i: ch for i, ch in enumerate(all_words)}
        pickle.dump(char_to_id, open(self.config.char_to_id_file, 'wb'))
        pickle.dump(id_to_char, open(self.config.id_to_char_file, 'wb'))

        # vocab_vectors:
        self.vocab_vectors = []
        for caption in self.captions:
            temp_vector = [char_to_id[word.lower()] for word in caption.split(' ')]
            self.vocab_vectors.append(temp_vector)
        pickle.dump(self.vocab_vectors, open(self.config.vocab_vectors_file, 'wb'))

        # train_images:
        train_images = set(open(self.config.train_images_file, 'r').read().strip().split('\n'))
        self.train_images, self.train_captions = self.filter_image_vector(train_images)

        val_images = set(open(self.config.val_images_file, 'r').read().strip().split('\n'))
        self.val_images, self.val_captions = self.filter_image_vector(val_images)

        test_images = set(open(self.config.test_images_file, 'r').read().strip().split('\n'))
        self.test_images, self.test_captions = self.filter_image_vector(test_images)

    def filter_image_vector(self, images):
        """产生对应文件夹中的image feature vector，和对应的caption vector"""
        img_features = np.load(self.config.image_path_to_vector_file)
        temp_image_vectors = []
        temp_caption_vectors = []
        for idx, img_id in enumerate(self.images_names):
            if img_id in images:
                temp_image_vectors.append(img_features[()][self.config.image_folder + img_id])
                temp_caption_vectors.append(self.vocab_vectors[idx])
        return temp_image_vectors, temp_caption_vectors

    def batch_iter(self):
        """产生batches"""
        length = len(self.train_images)
        num_batches_per_epoch = length // self.training_config.batch_size
        for epoch in range(self.training_config.epochs):
            for batch_num in range(num_batches_per_epoch):
                start_index = batch_num * self.training_config.batch_size
                end_index = start_index + self.training_config.batch_size
                img_batch = self.train_images[start_index:end_index]

                batch = self.train_captions[start_index:end_index]
                seq_len = max(map(len, batch))
                # TODO: unk
                x_batch = np.full([self.training_config.batch_size, seq_len], 9000, dtype=np.int32)

                seq_len = []
                for row in range(self.training_config.batch_size):
                    row_len = len(batch[row])
                    seq_len.append(row_len)
                    x_batch[row, :row_len] = batch[row]

                y_batch = np.copy(x_batch)
                y_batch[:, :-1] = x_batch[:, 1:]

                yield img_batch, x_batch, y_batch, seq_len


if __name__ == '__main__':
    data = Data()
    batches = data.batch_iter()
    for batch in batches:
        img_feature, x, y, seq_len = batch
        # img_feature = np.reshape(img_feature, np.shape(img_feature)[1])
        print(len(img_feature), 'img_feature:', np.shape(img_feature))
        print('img_feature:', img_feature)
        print(len(x), 'x:', x)
        print(len(y), 'y:', y)
